import os
import json
import numpy as np
from scipy.optimize import linear_sum_assignment

# 类别映射
category_mapping = {
    "car": 0, "forklift": 1, "tractor": 2, "truck": 3, "flatbed": 4,
    "trailer": 5, "bicycle": 6, "cyclist": 7, "pedestrian": 8,
    "traffic_cone": 9, "water_horse": 10, "guardrail": 11,
}
UNKNOWN_LABEL = -1
class_names = list(category_mapping.keys())
num_classes = len(class_names)
PEDESTRIAN_IDX = category_mapping["pedestrian"]


def parse_gt_label(json_path):
    with open(json_path, 'r') as f:
        data = json.load(f)
    boxes, labels = [], []
    for obj in data:
        cls = obj.get("obj_type")
        if cls not in category_mapping:
            continue
        pos = obj["psr"]["position"]
        size = obj["psr"]["scale"]
        x, y, z = pos["x"], pos["y"], pos["z"]
        l, w, h = size["y"], size["x"], size["z"]
        boxes.append([x, y, z, l, w, h])
        labels.append(category_mapping[cls])
    return np.array(boxes), labels


def parse_pred_txt(txt_path):
    boxes, labels = [], []
    with open(txt_path, 'r') as f:
        for line in f:
            parts = line.strip().split(',')
            if len(parts) != 8:
                continue
            cls = parts[0]
            if cls.lower() == 'unknown':
                label = UNKNOWN_LABEL
            elif cls in category_mapping:
                label = category_mapping[cls]
            else:
                continue
            x, y, z = map(float, parts[1:4])
            l, w, h = map(float, parts[4:7])
            boxes.append([x, y, z, l, w, h])
            labels.append(label)
    return np.array(boxes), labels


def calculate_cost_matrix(b1, b2):
    m, n = len(b1), len(b2)
    cost = np.zeros((m, n))
    for i in range(m):
        for j in range(n):
            cost[i, j] = np.linalg.norm(b1[i][:3] - b2[j][:3])
    return cost


def match_3d_boxes(gt_boxes, pred_boxes):
    cost = calculate_cost_matrix(gt_boxes, pred_boxes)
    return linear_sum_assignment(cost)

# 初始化路径和统计结构
gt_dir = 'dataset/houronghe/houronghe1/label'
pred_dir = 'dataset/houronghe/detect1'
distance_ranges = [10, 20, 30]
global_stats = [dict(GT=0, TP=0, FP=0, FN=0) for _ in range(num_classes)]
global_dist = {d: [dict(GT=0, TP=0, FP=0, FN=0) for _ in range(num_classes)] for d in distance_ranges}

# 用于记录出现行人误报的帧
ped_fp_frames = []

# 遍历每一帧
for file in os.listdir(gt_dir):
    if not file.endswith('.json'):
        continue
    tok = os.path.splitext(file)[0]
    gt_boxes, gt_labels = parse_gt_label(os.path.join(gt_dir, file))
    pred_boxes, pred_labels = parse_pred_txt(os.path.join(pred_dir, tok+".txt"))

    # 初始化 FN = GT
    for i, cls in enumerate(gt_labels):
        global_stats[cls]['GT'] += 1
        global_stats[cls]['FN'] += 1
        d_gt = np.linalg.norm(gt_boxes[i][:3])
        for d in distance_ranges:
            if d_gt <= d:
                global_dist[d][cls]['GT'] += 1
                global_dist[d][cls]['FN'] += 1

    # 若无预测，跳过后续
    if len(pred_boxes) == 0:
        continue

    # 匹配
    rows, cols = match_3d_boxes(gt_boxes, pred_boxes)
    matched_gt, matched_pr = set(rows), set(cols)

    # 本帧是否出现行人误报
    ped_fp_in_this_frame = False

    # 处理匹配对
    for gt_i, pr_i in zip(rows, cols):
        dist = np.linalg.norm(gt_boxes[gt_i][:3] - pred_boxes[pr_i][:3])
        cls_gt = gt_labels[gt_i]
        cls_pr = pred_labels[pr_i]
        gt_dist = np.linalg.norm(gt_boxes[gt_i][:3])

        if dist <= 4.0:
            if cls_pr == cls_gt:
                # 真阳性
                global_stats[cls_gt]['TP'] += 1
                global_stats[cls_gt]['FN'] -= 1
                for d in distance_ranges:
                    if gt_dist <= d:
                        global_dist[d][cls_gt]['TP'] += 1
                        global_dist[d][cls_gt]['FN'] -= 1
            elif cls_pr == UNKNOWN_LABEL:
                # Unknown 匹配，只减 FN
                global_stats[cls_gt]['FN'] -= 1
                for d in distance_ranges:
                    if gt_dist <= d:
                        global_dist[d][cls_gt]['FN'] -= 1
            else:
                # 错误匹配，当作 FP
                global_stats[cls_pr]['FP'] += 1
                # 如果错误匹配为行人，也记录本帧
                if cls_pr == PEDESTRIAN_IDX:
                    ped_fp_in_this_frame = True
                for d in distance_ranges:
                    if gt_dist <= d:
                        global_dist[d][cls_pr]['FP'] += 1

    # 剩余预测，作为 FP；并检测行人 FP
    for pr_i, cls in enumerate(pred_labels):
        if pr_i not in matched_pr:
            if cls >= 0:
                global_stats[cls]['FP'] += 1
                d_pr = np.linalg.norm(pred_boxes[pr_i][:3])
                for d in distance_ranges:
                    if d_pr <= d:
                        global_dist[d][cls]['FP'] += 1
                # 如果是行人误报，记录本帧
                if cls == PEDESTRIAN_IDX:
                    ped_fp_in_this_frame = True
    if ped_fp_in_this_frame:
        ped_fp_frames.append(tok)

# 打印结果
print('=== 全局 Overall (All Distances) ===')
print(f"{'class':<12}{'GT':>6}{'TP':>6}{'FP':>6}{'FN':>6}{'Prec':>8}{'Rec':>8}")
for idx, name in enumerate(class_names):
    G = global_stats[idx]['GT']
    T = global_stats[idx]['TP']
    Fp = global_stats[idx]['FP']
    Fn = global_stats[idx]['FN']
    P = T / (T + Fp + 1e-6)
    R = T / (T + Fn + 1e-6)
    print(f"{name:<12}{G:6}{T:6}{Fp:6}{Fn:6}{P:8.3f}{R:8.3f}")

for d in distance_ranges:
    print(f"\n=== Distance ≤ {d} m ===")
    print(f"{'class':<12}{'GT':>6}{'TP':>6}{'FP':>6}{'FN':>6}{'Prec':>8}{'Rec':>8}")
    for idx, name in enumerate(class_names):
        G = global_dist[d][idx]['GT']
        T = global_dist[d][idx]['TP']
        Fp = global_dist[d][idx]['FP']
        Fn = global_dist[d][idx]['FN']
        P = T / (T + Fp + 1e-6)
        R = T / (T + Fn + 1e-6)
        print(f"{name:<12}{G:6}{T:6}{Fp:6}{Fn:6}{P:8.3f}{R:8.3f}")

# 最后，输出所有包含行人误报的帧名称
if ped_fp_frames:
    print("\n=== 含行人误报的帧列表 ===")
    for tok in ped_fp_frames:
        print(tok)
else:
    print("\n未检测到任何行人误报的帧。")
